from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
from fastapi.responses import FileResponse
import io
import sys
import os
from dotenv import load_dotenv
sys.path.append(os.path.dirname(__file__))

from .app.model_npy import predict  # ← 本地使用相对路径导入
from llm_agent import generate_medical_advice  # ← 新增LangChain集成

# 加载环境变量
load_dotenv()

app = FastAPI(
    title="Brain Diagnosis API",
    description="脑部影像分类与AI就医建议系统",
    version="2.0.0"
)

from fastapi.middleware.cors import CORSMiddleware

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

@app.get("/.well-known/ai-plugin.json", include_in_schema=False)
def serve_ai_plugin():
    file_path = os.path.join(os.path.dirname(__file__), ".well-known", "ai-plugin.json")
    return FileResponse(path=file_path, media_type="application/json")

@app.post("/classify")
async def classify_npy(file: UploadFile = File(...)):
    """
    接收 .npy 文件作为输入图像，返回预测结果和AI就医建议
    默认从 batch 中取第 0 张图像进行推理
    """
    try:
        contents = await file.read()
        array = np.load(io.BytesIO(contents))

        if array.ndim not in [3, 4]:
            return JSONResponse(
                status_code=400,
                content={"error": f"输入维度不合法，应为 (H,W,C) 或 (B,H,W,C)，当前为 {array.shape}"}
            )

        # 获取原始预测结果
        result = predict(array)
        
        # 生成AI就医建议
        advice = generate_medical_advice({
            "label": result["label"],
            "confidence": result["confidence"],
            "probabilities": result["probabilities"]
        })
        
        # 合并结果
        result["advice"] = advice
        
        return JSONResponse(content=result)

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": f"处理失败: {str(e)}"}
        )

@app.post("/analyze")
async def analyze_result(prediction: dict):
    """
    独立接口：根据已有预测结果生成就医建议
    请求体示例：
    {
        "label": "AD",
        "confidence": 0.95,
        "probabilities": {"0": 0.05, "1": 0.95}
    }
    """
    try:
        advice = generate_medical_advice(prediction)
        return JSONResponse(content={"advice": advice})
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": f"建议生成失败: {str(e)}"}
        )

# 允许直接运行本地开发服务器
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)